"""
内容：获取中间层的输出
日期：2020年7月8日
作者：Howie
"""

import numpy as np
from keras.models import Sequential, Model
from keras.layers import Dense, MaxPooling2D, Conv2D, Flatten
import matplotlib.pyplot as plt

DATASET_PATH = '../dataset/mnist/mnist.npz'  # 数据集路径
# 超参
N_SAMPLES_TRAIN = 60000     # 训练集样本数量
N_SAMPLES_TEST = 10000  # 测试集样本数量
IMG_SIZE = 28
N_CLS = 10
EPOCH = 2
BATCH_SIZE = 64


class MLP:
    def __init__(self):
        self.model = Sequential()
        self.model.add(Dense(units=5, input_dim=1, activation='relu'))
        self.model.add(Dense(units=1))


class CNN:
    def __init__(self):
        self.model = Sequential()
        self.model.add(
            Conv2D(
                32, (3, 3), activation='relu', input_shape=(
                    IMG_SIZE, IMG_SIZE, 1), name='Conv1'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Conv2D(32, (3, 3), activation='relu', name='Conv2'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Conv2D(32, (3, 3), activation='relu', name='Conv3'))
        self.model.add(MaxPooling2D(pool_size=(2, 2)))
        self.model.add(Flatten())
        self.model.add(Dense(256, activation='relu'))
        self.model.add(Dense(N_CLS // 2, activation='sigmoid'))

    def conv_output(self, layer_name, img):
        """
        # 获取中间层的输出
        :param layer_name: 层名称
        :param img: 输入图片
        :return: 特征图
        """
        input_img = self.model.input

        try:
            # 输出特征图
            out_conv = self.model.get_layer(layer_name).output
        except:
            raise Exception('Not layer named {}!'.format(layer_name))

        # 创建一个新的模型来输出你所感兴趣的层
        intermediate_layer_model = Model(inputs=input_img, outputs=out_conv)
        intermediate_output = intermediate_layer_model.predict(img)

        return intermediate_output[0]

if __name__ == '__main__':
    with np.load(DATASET_PATH) as data:
        X = data['x_train'][0]
    X = X.reshape(1, IMG_SIZE, IMG_SIZE, 1).astype('float32') / 255.0
    model = CNN()
    plt.figure(0)
    plt.subplot(2, 2, 1)
    plt.title('Origin')
    plt.xticks([])
    plt.yticks([])
    plt.imshow(X[0, :, :, 0])
    for i in range(3):
        layer_name = 'Conv' + str(i+1)
        intermediate_conv_output = model.conv_output(layer_name, X)
        plt.subplot(2, 2, i+2)
        plt.title(layer_name)
        plt.xticks([])
        plt.yticks([])
        plt.imshow(intermediate_conv_output[:, :, 0])
    plt.savefig('./logs/Intermediate Conv Output.pdf')
    plt.show()

